/**
 * Copyright (c) 2018, Alex. All rights reserved.
 */
package com.oschina.bottle.coder.asymmetric;

import com.oschina.bottle.coder.CoderException;
import com.oschina.bottle.common.constants.Charset;
import java.security.Key;
import java.security.KeyFactory;
import java.security.KeyPair;
import java.security.KeyPairGenerator;
import java.security.PrivateKey;
import java.security.PublicKey;
import java.security.SecureRandom;
import java.security.Signature;
import java.security.spec.KeySpec;
import java.security.spec.PKCS8EncodedKeySpec;
import java.security.spec.X509EncodedKeySpec;
import java.util.Base64;
import javax.crypto.Cipher;
import org.apache.commons.lang3.ArrayUtils;

/**
 * @author <a href="mailto:zhongchao@gegejia.com">zhong</a>
 * @version 1.0 2018/1/4
 * @since 1.0
 */
public class AsymmetricCoder {

    public static final String SIGNATURE_ALGORITHM = "MD5withRSA";

    private AsymmetricAlgorithm algorithm;

    /**
     * 最大加密明文大小
     */
    private static final int MAX_ENCRYPT_BLOCK = 117;

    /**
     * 最大解密密文大小
     */
    private static final int MAX_DECRYPT_BLOCK = 128;

    /**
     * 私钥
     */
    private String privateKey;

    /**
     * 公钥
     */
    private String publicKey;

    public KeyData genKeyPair() {

        try {

            /** RSA算法要求有一个可信任的随机数源 */
            SecureRandom secureRandom = new SecureRandom();

            /** 为RSA算法创建一个KeyPairGenerator对象 */
            KeyPairGenerator keyPairGenerator = KeyPairGenerator.getInstance(algorithm.getName());

            /** 利用上面的随机数据源初始化这个KeyPairGenerator对象 */
            keyPairGenerator.initialize(1024, secureRandom);

            /** 生成密匙对 */
            KeyPair keyPair = keyPairGenerator.generateKeyPair();

            /** 得到公钥 */
            Key publicKey = keyPair.getPublic();

            /** 得到私钥 */
            Key privateKey = keyPair.getPrivate();

            byte[] privateKeyBytes = privateKey.getEncoded();
            byte[] publicKeyBytes = publicKey.getEncoded();

            return new KeyData(Base64.getEncoder().encodeToString(privateKeyBytes), Base64.getEncoder().encodeToString(publicKeyBytes));
        } catch (Exception ex) {
            throw new CoderException(ex, "生成密钥对失败");
        }
    }

    public AsymmetricCoder(AsymmetricAlgorithm algorithm, String privateKey, String publicKey) {
        this.algorithm = algorithm;
        this.privateKey = privateKey;
        this.publicKey = publicKey;
    }

    public byte[] encrypt(byte[] data, KeyType keyType) {
        return doFinal(data, keyType, Cipher.ENCRYPT_MODE);
    }

    public String encrypt(String data, KeyType keyType) {
        return Base64.getEncoder().encodeToString(encrypt(Charset.formatUTF8(data), keyType));
    }

    public byte[] decrypt(byte[] data, KeyType keyType) {
        return doFinal(data, keyType, Cipher.DECRYPT_MODE);
    }

    public String decrypt(String data, KeyType keyType) {
        return new String(decrypt(Base64.getDecoder().decode(data), keyType));
    }

    private byte[] doFinal(byte[] data, KeyType keyType, int mode) {
        byte[] dataReturn = null;
        try {
            Key key = toKey(keyType);
            Cipher cipher = initCipher();
            cipher.init(mode, key);

            // 加密时超过117字节就报错。为此采用分段加密的办法来加密
            int max = mode == Cipher.DECRYPT_MODE ? MAX_DECRYPT_BLOCK : MAX_ENCRYPT_BLOCK;
            for (int i = 0; i < data.length; i += max) {
                byte[] doFinal = cipher.doFinal(subarray(data, i, i + max));
                dataReturn = ArrayUtils.addAll(dataReturn, doFinal);
            }
        } catch (Exception e) {
            throw new CoderException(e, "%s加解密失败", algorithm.getName());
        }

        return dataReturn;
    }

    private Cipher initCipher() {
        try {
            return Cipher.getInstance(algorithm.getName());
        } catch (Exception e) {
            throw new CoderException("初始化cipher失败，算法：%s", algorithm.getName());
        }
    }

    private Key toKey(KeyType keyType) {

        try {
            String key = getKeyByType(keyType);
            KeySpec keySpec;
            KeyFactory keyFactory;
            switch (keyType) {
                case PRIVATE:
                    keySpec = new PKCS8EncodedKeySpec(Base64.getDecoder().decode(key));
                    keyFactory= KeyFactory.getInstance(algorithm.getName());
                    return keyFactory.generatePrivate(keySpec);
                case PUBLIC:
                    keySpec = new X509EncodedKeySpec(Base64.getDecoder().decode(key));
                    keyFactory= KeyFactory.getInstance(algorithm.getName());
                    return keyFactory.generatePublic(keySpec);
                default:
                    throw new CoderException("不支持的keyType类型：%s", keyType.name());
            }
        } catch (Exception ex) {
            throw new CoderException(ex, "初始化key出错");
        }
    }

    private String getKeyByType(KeyType keyType) {
        switch (keyType) {
            case PRIVATE:
                if (null == privateKey) {
                    throw new IllegalArgumentException("the privateKey must be null");
                }
                return privateKey;
            case PUBLIC:
                if (null == publicKey) {
                    throw new IllegalArgumentException("the publicKey must be null");
                }
                return publicKey;
            default:
                throw new CoderException("不支持的keyType类型：%s", keyType.name());
        }
    }

    /**
     * 签名
     *
     * @param data      明文数据
     * @return
     */
    public String sign(String data) {
        try {
            PrivateKey key = (PrivateKey) toKey(KeyType.PRIVATE);
            Signature signature = Signature.getInstance(SIGNATURE_ALGORITHM);
            signature.initSign(key);
            signature.update(Charset.formatUTF8(data));

            return Base64.getEncoder().encodeToString(signature.sign());
        } catch (Exception ex) {
            throw new CoderException(ex, "生成签名失败");
        }
    }

    /**
     * 校验签名
     *
     * @param data      明文数据
     * @param sign      签名数据
     * @return
     */
    public boolean verify(String data, String sign) {
        try {
            PublicKey key = (PublicKey) toKey(KeyType.PUBLIC);
            Signature signature = Signature.getInstance(SIGNATURE_ALGORITHM);
            signature.initVerify(key);
            signature.update(Charset.formatUTF8(data));

            return signature.verify(Base64.getDecoder().decode(sign));
        } catch (Exception ex) {
            throw new CoderException(ex, "签名校验失败");
        }
    }

    /**
     * <p>Produces a new <code>byte</code> array containing the elements between the start and end indices.</p> <p>
     * <p>The start index is inclusive, the end index exclusive. Null array input produces null output.</p>
     *
     * @param array the array
     * @param startIndexInclusive the starting index. Undervalue (&lt;0) is promoted to 0, overvalue (&gt;array.length)
     * results in an empty array.
     * @param endIndexExclusive elements up to endIndex-1 are present in the returned subarray. Undervalue (&lt;
     * startIndex) produces empty array, overvalue (&gt;array.length) is demoted to array length.
     * @return a new array containing the elements between the start and end indices.
     * @since 2.1
     */
    private byte[] subarray(byte[] array, int startIndexInclusive, int endIndexExclusive) {
        if (array == null) {
            return null;
        }
        if (startIndexInclusive < 0) {
            startIndexInclusive = 0;
        }
        if (endIndexExclusive > array.length) {
            endIndexExclusive = array.length;
        }
        int newSize = endIndexExclusive - startIndexInclusive;
        if (newSize <= 0) {
            return new byte[0];
        }

        byte[] subarray = new byte[newSize];
        System.arraycopy(array, startIndexInclusive, subarray, 0, newSize);
        return subarray;
    }

}
